from flask import Flask, request, jsonify
from flask_cors import CORS
import json
import os
import base64
from PIL import Image
from io import BytesIO
import shutil
import numpy as np
import requests

app = Flask(__name__)
app.secret_key = b'_5#y3332323xd'
CORS(app, origins='*')

img_read_path = './img'
img_save_path = './img_out'
txt_save_path = './txt_out'
img_size = (640, 640)

detect_cmd = 'python ./predict.py'


def save_base64_image(base64_str, file_path):
    img_data = base64.b64decode(base64_str)
    img = Image.open(BytesIO(img_data))
    img = img.resize(img_size, Image.Resampling.LANCZOS)

    img.save(file_path)


def get_detections():
    all_detections = []
    for filename in os.listdir(img_save_path):
        # 读取图像文件并转换为 base64 编码
        with open(os.path.join(img_save_path, filename), 'rb') as f:
            image_data = f.read()
            image_base64 = base64.b64encode(image_data).decode('utf-8')
            # 获取图片后缀名，添加类型（先输出看下）
            image_base64_prefix = "data:image/" + filename.split('.')[-1] + ";base64,"
            # 组装结果
            result = {
                'image_name': filename,
                'image_base64': image_base64_prefix + image_base64
            }
            all_detections.append(result)

    # 读取识别log日志文件
    folder_path = txt_save_path
    file_list = [f for f in os.listdir(folder_path) if f.endswith(".txt")]

    matrix_list = []

    for file_name in file_list:
        file_path = os.path.join(folder_path, file_name)
        matrix_data = np.loadtxt(file_path)  # 使用适当的方法读取矩阵数据
        # matrix_list.append(matrix_data.tolist())
        total_elements = matrix_data.size
        ones_count = np.count_nonzero(matrix_data == 1)
        ones_percentage = (ones_count / total_elements) * 100
        matrix_shape = matrix_data.shape
        matrix_list.append(
            f":图片尺寸： {matrix_shape}, 总像素数 ：{total_elements}, 单车停放区域所占像素数：{ones_count}, 单车区域占比：{ones_percentage:.2f}%")

    # 组装最终json结果
    final_res = {
        'details': all_detections,
        'logs': matrix_list
    }

    # 清除目录
    rmDir()

    return final_res


def mkDir():
    # 处理图片和结果文本路径
    if not os.path.exists(os.path.join(img_read_path)):
        os.makedirs(os.path.join(img_read_path))
    if not os.path.exists(os.path.join(img_save_path)):
        os.makedirs(os.path.join(img_save_path))
    if not os.path.exists(os.path.join(txt_save_path)):
        os.makedirs(os.path.join(txt_save_path))


def rmDir():
    # 清除目录
    if os.path.exists(img_read_path):
        shutil.rmtree(img_read_path)
    if os.path.exists(img_save_path):
        shutil.rmtree(img_save_path)
    if os.path.exists(txt_save_path):
        shutil.rmtree(txt_save_path)


def process_img(index, image_code):
    # 判断图片类型
    img_type, img_code = image_code.split(',')[0], image_code.split(',')[1]
    start = img_type.find('/') + 1
    end = img_type.find(';')
    mime_type = img_type[start:end]
    if mime_type not in ['png', 'jpg', 'jpeg']:
        return jsonify({"code": 500, "msg": "图片类型错误，必须为png或jpg!"})
    save_base64_image(img_code, os.path.join(img_read_path, str(index) + '.' + mime_type))


@app.route('/image-deepLabV3-detect', methods=['POST'])
def detect_image_obj():
    try:
        # 获取传入的参数
        data = json.loads(request.get_data())
        image_codes = data.get('imageCodes')
        # 处理图片和结果文本路径
        mkDir()
        for index, image_code in enumerate(image_codes):
            process_img(index, image_code)
        # 执行预测命令
        cmd = detect_cmd
        print('正在执行预测脚本：', cmd)
        os.system(cmd)
        # 整理输出结果并删除目录
        return jsonify(get_detections())
    # 异常参数处理
    except Exception:
        return jsonify({'code': '500', 'msg': 'error'})


def get_yolo_obj_loc_list(image_code):
    obj_loc_list = []
    # API接口URL
    api_url = "http://47.103.135.45:9528/image-yolov5-detect"
    # 参数
    yolo_data = {
        "imageCodes": [image_code]
    }
    # 发起POST请求
    response = requests.post(api_url, json=yolo_data)
    # 检查响应状态码
    if response.status_code == 200:
        # 解析响应内容
        response_data = response.json()
        details_list = response_data["details"]
        first_details = details_list[0]
        obj_loc_list = first_details.get('detections')

    return obj_loc_list


objs = ['motor', 'private_bike', 'shared_bike']


# 设置成只处理单张图片
@app.route('/judge_if_parking_rightly', methods=['POST'])
def judge_if_parking_rightly():
    # try:
    # 获取传入参数
    data = json.loads(request.get_data())
    image_code = data.get('imageCode')
    # 调用YOLO获取位置集合信息
    obj_loc_list = get_yolo_obj_loc_list(image_code)
    # 处理图片和结果文本路径
    mkDir()
    # 保存图片（resize）
    process_img(0, image_code)
    # 执行DeepLabV3预测命令
    cmd = detect_cmd
    print('正在执行预测脚本：', cmd)
    os.system(cmd)
    # 读取语义分割识别结果
    pr = np.loadtxt(txt_save_path + '/0.txt', dtype=int)
    # 遍历obj_loc_list，检索单车总数、在区域内的单车数目，并统计百分占比
    bike_in_zone = 0
    bike_sum = 0
    for obj in obj_loc_list:
        class_name = objs[obj['class_id']]
        if class_name == 'shared_bike':
            bike_sum += 1
            # 只要共享单车与单车停放区域有交集就认为其停在其中
            x = obj['x']
            y = obj['y']
            width = obj['width']
            height = obj['height']
            x_min = x - width / 2
            y_min = y - height / 2
            x_max = x + width / 2
            y_max = y + height / 2
            x_pixel_min = int(x_min * img_size[0])
            y_pixel_min = int(y_min * img_size[1])
            x_pixel_max = int(x_max * img_size[0])
            y_pixel_max = int(y_max * img_size[1])

            contains_1 = False
            for y in range(y_pixel_min, y_pixel_max + 1):
                for x in range(x_pixel_min, x_pixel_max + 1):
                    if y >= 0 and y <= 639 and x >= 0 and x <= 639 and pr[y][x] == 1:
                        contains_1 = True
                        break
                if contains_1:
                    break
            if contains_1:
                bike_in_zone += 1
    # 清除工作目录
    rmDir()

    return jsonify({"code": 200,
                    "msg": f":单车总数： {bike_sum}, 位于停车区域的单车数目 ：{bike_in_zone}, 单车摆放正确率：{100.0 * bike_in_zone / bike_sum:.2f}%"})


# 异常参数处理
# except Exception:
#     return jsonify({'code': '500', 'msg': 'error'})


def get_yolo_obj_img(image_code):
    obj_loc_img = ''
    # API接口URL
    api_url = "http://47.103.135.45:9528/image-yolov5-detect"
    # 参数
    yolo_data = {
        "imageCodes": [image_code]
    }
    # 发起POST请求
    response = requests.post(api_url, json=yolo_data)
    # 检查响应状态码
    if response.status_code == 200:
        # 解析响应内容
        response_data = response.json()
        details_list = response_data["details"]
        first_details = details_list[0]
        obj_loc_img = first_details.get('image_base64')

    return obj_loc_img


def blend_and_encode(pr, base64_image):
    orininal_h = img_size[0]
    orininal_w = img_size[1]

    colors = [(0, 0, 0), (0, 100, 0)]

    seg_img = np.reshape(np.array(colors, np.uint8)[np.reshape(pr, [-1])], [orininal_h, orininal_w, -1])
    image = Image.fromarray(np.uint8(seg_img))

    # Decode the base64-encoded image and create a PIL Image object
    decoded_image = base64.b64decode(base64_image.split(',')[1])
    old_img = Image.open(BytesIO(decoded_image))

    # Blend the images
    blended_image = Image.blend(old_img, image, 0.5)
    # blended_image.show()

    # Convert the blended image to base64 with data prefix
    buffered = BytesIO()
    blended_image.save(buffered, format="JPEG")
    blended_image_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")

    return f"data:image/jpeg;base64,{blended_image_base64}"


# 设置成只处理单张图片
@app.route('/get_blend_detect_img', methods=['POST'])
def get_blend_img():
    try:
        # 获取传入参数
        data = json.loads(request.get_data())
        image_code = data.get('imageCode')
        # 调用YOLO获取图像编码
        obj_img_code = get_yolo_obj_img(image_code)
        # 处理图片和结果文本路径
        mkDir()
        # 保存图片（resize）
        process_img(0, image_code)
        # 执行DeepLabV3预测命令
        cmd = detect_cmd
        print('正在执行预测脚本：', cmd)
        os.system(cmd)
        # 读取语义分割识别结果
        pr = np.loadtxt(txt_save_path + '/0.txt', dtype=int)
        # 获取blend_img
        blend_img_code = blend_and_encode(pr, obj_img_code)
        # 清楚工作目录
        rmDir()

        return jsonify({'code': '200', 'msg': blend_img_code})
    # 异常参数处理
    except Exception:
        return jsonify({'code': '500', 'msg': 'error'})


if __name__ == '__main__':
    os.chdir(os.path.abspath(os.path.dirname(__file__)))
    app.run(debug=True, host='0.0.0.0', port='9529')
